import os
from typing import List, Dict, Optional, Callable

import cv2
import fitz
import numpy as np
from PIL import Image
from skimage.feature import local_binary_pattern


class HandDrawingExtractor:
    def __init__(self,
                 min_area_ratio: float = 0.01,
                 drawing_threshold: float = 0.6,
                 dpi: int = 300,
                 morph_kernel_size: int = 5,
                 edge_density_weight: float = 0.4,
                 line_score_weight: float = 0.3,
                 texture_weight: float = 0.1):
        self.min_area_ratio = min_area_ratio
        self.drawing_threshold = drawing_threshold
        self.dpi = dpi
        self.morph_kernel = np.ones((morph_kernel_size, morph_kernel_size), np.uint8)

        # 验证权重总和为1.0
        total_weight = edge_density_weight + line_score_weight + 0.2 + texture_weight
        if not np.isclose(total_weight, 1.0):
            raise ValueError("!=1.0")

        self.edge_density_weight = edge_density_weight
        self.line_score_weight = line_score_weight
        self.texture_weight = texture_weight

    def extract_drawings(self,
                         pdf_path: str,
                         output_dir: str = "output",
                         progress_callback: Optional[Callable[[float], None]] = None) -> List[Dict]:

        if not os.path.exists(pdf_path):
            raise FileNotFoundError(f"PDF文件不存在: {pdf_path}")

        os.makedirs(output_dir, exist_ok=True)
        doc = fitz.open(pdf_path)
        total_pages = len(doc)
        drawings = []

        for page_num in range(total_pages):
            if progress_callback:
                progress_callback(page_num / total_pages)

            page = doc.load_page(page_num)
            print(f"第 {page_num + 1} 张")
            try:
                pix = page.get_pixmap(dpi=self.dpi)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                cv_img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

                page_drawings = self._detect_drawings(cv_img)

                for i, drawing in enumerate(page_drawings):
                    x, y, w, h = drawing["position"]
                    roi = drawing["image"]
                    save_path = os.path.join(
                        output_dir,
                        f"page_{page_num + 1}hand_drawing_{i + 1}.png"
                    )
                    print(f"图片尺寸: {roi.size}")
                    roi.save(save_path)

                    drawings.append({
                        "page": page_num + 1,
                        "position": [x, y, w, h],
                        "image_path": save_path,
                        "image": roi
                    })
            except Exception as e:
                print(f"处理第 {page_num + 1} 页时出错: {str(e)}")
                continue

        doc.close()
        if progress_callback:
            progress_callback(1.0)
        return drawings

    def _detect_drawings(self, cv_img: np.ndarray, expand_pixels: int = 1000) -> List[Dict]:
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        gray = cv2.equalizeHist(gray)
        binary = cv2.adaptiveThreshold(
            gray, 255,
            cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY_INV, 21, 10
        )
        processed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, self.morph_kernel)

        img_h, img_w = gray.shape
        min_area = img_h * img_w * self.min_area_ratio
        drawings = []
        contours, _ = cv2.findContours(processed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area >= min_area:
                x, y, w, h = cv2.boundingRect(cnt)

                x1 = max(0, x - expand_pixels)  # 向左
                y1 = max(0, y - 500)  # 向上
                x2 = min(img_w, x + w + expand_pixels)  # 向右
                y2 = min(img_h, y + h + expand_pixels)  # 向下
                expanded_w = x2 - x1
                expanded_h = y2 - y1
                expanded_roi = cv_img[y1:y2, x1:x2]
                if self._is_hand_drawing(expanded_roi):
                    drawings.append({
                        "position": [x1, y1, expanded_w, expanded_h],  # 坐标 高宽
                        "image": Image.fromarray(cv2.cvtColor(expanded_roi, cv2.COLOR_BGR2RGB)),
                        "original_bbox": [x, y, w, h]
                    })

        return drawings

    def _is_hand_drawing(self, roi: np.ndarray) -> bool:

        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)

        #  边缘密度检测
        edges = cv2.Canny(gray, 50, 150)
        edge_density = np.sum(edges > 0) / edges.size

        # 线条曲率分析
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=20, minLineLength=20, maxLineGap=5)
        line_score = len(lines) / 100 if lines is not None else 0

        # 纹理分析（新增）
        texture_score = self._texture_analysis(roi)

        # 综合评分（使用加权平均）
        drawing_score = (
                edge_density * self.edge_density_weight +
                line_score * self.line_score_weight +
                texture_score * self.texture_weight
        )

        return drawing_score >= self.drawing_threshold

    def _texture_analysis(self, roi: np.ndarray) -> float:
        """纹理分析（LBP方法）"""
        gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        lbp = local_binary_pattern(gray, 8, 1, method='uniform')
        hist, _ = np.histogram(lbp, bins=20, range=(0, 20))
        texture_score = np.std(hist) / np.mean(hist)
        return min(1.0, max(0.0, texture_score))


if __name__ == "__main__":
    def print_progress(percentage: float):
        print(f"处理进度: {percentage:.1%}", end="", flush=True)


    try:
        extractor = HandDrawingExtractor(
            min_area_ratio=0.006,  # 检测区域
            drawing_threshold=0.55,  # 阈值
            dpi=300,  # 分辨率
            morph_kernel_size=7,  # 形态学核
            edge_density_weight=0.3,
            line_score_weight=0.3,
            texture_weight=0.2
        )

        pdf_path = r"pdf\US-20040180768-A1.pdf"
        output_dir = f"extracted_drawings//{os.path.basename(pdf_path).split('.')[0]}"

        print("开始")
        drawings = extractor.extract_drawings(
            pdf_path,
            output_dir=output_dir,
            progress_callback=print_progress
        )

        print(f"结束 {len(drawings)} ")

    except Exception as e:
        print(f"error: {str(e)}")
